#! python
#-*- coding: utf-8 -*- 

from abc import ABC, abstractmethod
from typing import Tuple, Callable, Optional, TypeAlias
from numpy.typing import NDArray
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import font_manager
from matplotlib.axes import Axes
from .distribution import Real, Distribution

Statistics: TypeAlias = Callable[[NDArray[Real]], Real]

class Sampler(ABC):
    """采样器的抽象基类"""
    def __init__(self, chinese: bool = True):
        super().__init__()
        self._chinese = chinese
        self.setup()

    @abstractmethod
    def sample(self, p: Distribution, count: int) -> Tuple[NDArray[Real], NDArray[Real]]:
        pass
    
    def setup(self):
        """ 设置中文字体 """
        try:
            font_path = "C:/Windows/Fonts/simhei.ttf"
            font_prop = font_manager.FontProperties(fname=font_path)
            plt.rcParams['font.sans-serif'] = [font_prop.get_name()]
            plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
            self._chinese = True
        except:
            print("警告: 中文字体设置失败，图表可能无法正常显示中文")
            self._chinese = False
    
    def translate(self, chinese: str, english: str) -> str:
        if self._chinese:
            return chinese
        else:
            return english

    def expectation(
        self,
        statistics: Statistics,
        p: Distribution,
        num_samples: int = 10000,
        plot: bool = False,
        ax: Optional[Axes] = None,
    ) -> Tuple[Real, dict]:
        """ 评估统计量在该分布下的数学期望 
        
        参数:
            statistics: 统计量函数，接受样本数组返回标量值
            p: 目标概率分布
            num_samples: 采样数量，默认10000
            plot: 是否绘图，默认False
            ax: 可选的matplotlib Axes对象
            
        返回:
            (average_value, context) 统计量的期望值和包含rms等的上下文字典
        """
        data, log_weights = self.sample(p, num_samples)
        max_log = np.max(log_weights)
        weights = np.exp(log_weights - max_log)  # 数值稳定处理
        
        # 计算基本统计量
        value = statistics(data)
        weighted_value = value * weights
        weighted_square = value * value * weights
        
        # 计算各阶矩
        mu_A = np.sum(weighted_value) / num_samples       
        mu_B = np.sum(weights) / num_samples                 
        mu_fw2 = np.sum(weighted_value * weights) / num_samples
        mu_f2w2 = np.sum(weighted_square * weights) / num_samples
        mu_w2 = np.sum(weights * weights) / num_samples
        mu_B2 = mu_B * mu_B
        mu_B3 = mu_B2 * mu_B
        mu_B4 = mu_B3 * mu_B
        # 计算方差
        var = (
            mu_f2w2 * mu_B2
            - 2 * mu_A * mu_fw2 / mu_B3
            + mu_A * mu_A * mu_w2 / mu_B4
        ) / num_samples
        
        # 计算标准差
        rms = np.sqrt(var) if var > 0 else 0.0
        
        context = {
            "variance": var,
            "standard_error": rms,
            "average_value": mu_A,
            "average_weights": mu_B,
            "average_weighted_value": mu_fw2,
            "average_weighted_square": mu_f2w2,
            "average_weighted_square_weights": mu_w2,
            "average_weights_squared": mu_B2,
            "average_weights_cubed": mu_B3,
            "average_weights_quartic": mu_B4,
            "num_samples": num_samples    
        }       
        return mu_A, context


class DirectSampler(Sampler):
    """直接蒙特卡洛采样"""
    def __init__(self):
        super().__init__()
        
    def sample(self, p: Distribution, count: int) -> Tuple[NDArray[Real], NDArray[Real]]:
        """执行直接蒙特卡洛采样
        
        参数:
            count: 采样数量
            
        返回:
            元组(data, log_weights)，其中:
            - data: 采样数据，形状为(dimension, count)
            - log_weights: 对数权重数组，形状为(1, count)，所有元素为0.0
            
        示例:
        >>> p = RealGaussian()
        >>> sampler = DirectSampler()
        >>> data, log_weights = sampler.sample(p, 100)
        >>> data.shape[1] == 100  # 检查采样数量
        True
        >>> log_weights.shape == (1, 100)  # 检查权重形状
        True
        >>> np.allclose(log_weights, 0.0)  # 检查权重值
        True
        >>> isinstance(data[0,0], Real)  # 检查数据类型
        True
        """
        data = p.random(count)
        shape = (1, data.shape[1])
        log_weights = np.zeros(shape)
        return data, log_weights

class ImportanceSampler(Sampler):
    """重要性采样"""
    def __init__(self, proposal: Distribution):
        super().__init__()
        self._proposal = proposal

    def sample(self, p: Distribution, count: int) -> Tuple[NDArray[Real], NDArray[Real]]:
        """执行重要性采样
        
        参数:
            p: 目标概率分布
            count: 采样数量
            
        返回:
            元组(data, log_weights)，其中:
            - data: 采样数据，形状为(dimension, count)
            - log_weights: 对数权重数组，形状为(1, count)
            
        示例:
        >>> # 基本用法测试
        >>> target = RealGaussian()
        >>> proposal = RealGaussian()
        >>> sampler = ImportanceSampler(proposal)
        >>> data, log_weights = sampler.sample(target, 100)
        >>> data.shape[1] == 100  # 检查采样数量
        True
        >>> log_weights.shape == (1, 100)  # 检查权重形状
        True
        >>> isinstance(data[0,0], Real)  # 检查数据类型
        True
        
        >>> # 权重计算测试
        >>> target = RealGaussian()
        >>> proposal = RealGaussian()
        >>> sampler = ImportanceSampler(proposal)
        >>> data, log_weights = sampler.sample(target, 100)
        >>> logp = target.logp(data)
        >>> logq = proposal.logp(data)
        >>> np.allclose(log_weights, logp - logq)  # 验证权重计算
        True
        
        >>> # 期望值估计测试
        >>> target = RealGaussian()
        >>> proposal = RealGaussian()
        >>> sampler = ImportanceSampler(proposal)
        >>> direct_sampler = DirectSampler()
        >>> # 重要性采样估计
        >>> imp_mean, imp_rms = sampler.expectation(lambda x: x, target, 10000)
        >>> # 直接采样估计
        >>> dir_mean, dir_rms = direct_sampler.expectation(lambda x: x, target, 10000)
        >>> # 比较两种方法的估计结果
        >>> np.all(abs(imp_mean - dir_mean) < 0.1)  # 均值估计应接近
        np.True_
        >>> np.all(abs(imp_rms - dir_rms) < 0.1)  # 标准差估计应接近
        np.True_
        """
        data = self._proposal.random(count)
        logp = p.logp(data)
        logq = self._proposal.logp(data)
        log_weights = (logp - logq).reshape(1, -1)  # 确保形状为(1, count)
        return data, log_weights


if __name__ == "__main__":
    import doctest
    doctest.testmod(optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, verbose=True)